"""
动态调度配置模块
从数据库读取和更新调度规则
"""
import os
from typing import Dict, Any, Optional
from datetime import timedelta
from celery.schedules import crontab
import motor.motor_asyncio
from pydantic import BaseModel, Field
import logging

logger = logging.getLogger(__name__)

class ScheduleConfig(BaseModel):
    """调度配置模型"""
    name: str = Field(..., description="任务名称")
    task: str = Field(..., description="任务函数路径")
    schedule_type: str = Field(..., description="调度类型: interval/crontab")
    schedule_value: Dict[str, Any] = Field(..., description="调度参数")
    enabled: bool = Field(True, description="是否启用")
    queue: str = Field('normal_priority', description="任务队列")
    priority: int = Field(5, description="优先级")
    expires: Optional[int] = Field(None, description="过期时间（秒）")
    args: list = Field(default_factory=list, description="任务参数")
    kwargs: dict = Field(default_factory=dict, description="任务关键字参数")

class DynamicScheduler:
    """动态调度器"""
    
    def __init__(self):
        self.mongo_uri = os.getenv('MONGODB_URI', 'mongodb://localhost:27017')
        self.db_name = os.getenv('MONGODB_DB', 'crawler_scheduler')
        self.client = None
        self.db = None
        
    async def connect(self):
        """连接MongoDB"""
        self.client = motor.motor_asyncio.AsyncIOMotorClient(self.mongo_uri)
        self.db = self.client[self.db_name]
        logger.info("Connected to MongoDB for dynamic scheduling")
        
    async def close(self):
        """关闭连接"""
        if self.client:
            self.client.close()
            
    async def get_schedules(self) -> Dict[str, Any]:
        """
        获取所有启用的调度配置
        
        Returns:
            调度配置字典
        """
        schedules = {}
        collection = self.db.schedule_configs
        
        async for doc in collection.find({'enabled': True}):
            config = ScheduleConfig(**doc)
            
            # 转换调度类型
            if config.schedule_type == 'interval':
                schedule = timedelta(**config.schedule_value)
            elif config.schedule_type == 'crontab':
                schedule = crontab(**config.schedule_value)
            else:
                logger.warning(f"Unknown schedule type: {config.schedule_type}")
                continue
                
            # 构建调度配置
            schedule_config = {
                'task': config.task,
                'schedule': schedule,
                'args': config.args,
                'kwargs': config.kwargs,
                'options': {
                    'queue': config.queue,
                    'priority': config.priority,
                }
            }
            
            if config.expires:
                schedule_config['options']['expires'] = config.expires
                
            schedules[config.name] = schedule_config
            
        logger.info(f"Loaded {len(schedules)} dynamic schedules")
        return schedules
        
    async def add_schedule(self, config: ScheduleConfig) -> bool:
        """
        添加新的调度配置
        
        Args:
            config: 调度配置
            
        Returns:
            是否成功
        """
        collection = self.db.schedule_configs
        
        try:
            await collection.insert_one(config.dict())
            logger.info(f"Added schedule: {config.name}")
            return True
        except Exception as e:
            logger.error(f"Failed to add schedule: {e}")
            return False
            
    async def update_schedule(self, name: str, updates: Dict[str, Any]) -> bool:
        """
        更新调度配置
        
        Args:
            name: 任务名称
            updates: 更新内容
            
        Returns:
            是否成功
        """
        collection = self.db.schedule_configs
        
        try:
            result = await collection.update_one(
                {'name': name},
                {'$set': updates}
            )
            if result.modified_count > 0:
                logger.info(f"Updated schedule: {name}")
                return True
            return False
        except Exception as e:
            logger.error(f"Failed to update schedule: {e}")
            return False
            
    async def delete_schedule(self, name: str) -> bool:
        """
        删除调度配置
        
        Args:
            name: 任务名称
            
        Returns:
            是否成功
        """
        collection = self.db.schedule_configs
        
        try:
            result = await collection.delete_one({'name': name})
            if result.deleted_count > 0:
                logger.info(f"Deleted schedule: {name}")
                return True
            return False
        except Exception as e:
            logger.error(f"Failed to delete schedule: {e}")
            return False
            
    async def toggle_schedule(self, name: str, enabled: bool) -> bool:
        """
        启用/禁用调度
        
        Args:
            name: 任务名称
            enabled: 是否启用
            
        Returns:
            是否成功
        """
        return await self.update_schedule(name, {'enabled': enabled})
        
    async def init_default_schedules(self):
        """初始化默认调度配置"""
        default_schedules = [
            ScheduleConfig(
                name='crawl-weibo-hot',
                task='src.tasks.crawler.crawl_weibo',
                schedule_type='interval',
                schedule_value={'minutes': 15},
                queue='crawler_queue',
                priority=8,
                expires=600
            ),
            ScheduleConfig(
                name='crawl-zhihu-hot',
                task='src.tasks.crawler.crawl_zhihu',
                schedule_type='interval',
                schedule_value={'minutes': 30},
                queue='crawler_queue',
                priority=7,
                expires=1200
            ),
            ScheduleConfig(
                name='health-check',
                task='src.tasks.health.health_check',
                schedule_type='interval',
                schedule_value={'minutes': 5},
                queue='health_queue',
                priority=9,
                expires=240
            ),
        ]
        
        collection = self.db.schedule_configs
        
        for schedule in default_schedules:
            # 检查是否已存在
            existing = await collection.find_one({'name': schedule.name})
            if not existing:
                await self.add_schedule(schedule)
                
        logger.info("Initialized default schedules")